#!/usr/bin/env python

import os
import argparse
import numpy as np
from astropy.io import fits
from astropy.modeling.fitting import LevMarLSQFitter
from astropy.modeling.functional_models import Gaussian2D
from scipy.ndimage import binary_erosion, binary_dilation
from scipy.stats import sigmaclip
from skimage.morphology import label
import matplotlib.pyplot as plt
import sys


def measure_beam(fdata, fbias=None, limit=0.1, sigma=None,
                 xcent=None, ycent=None, xsize=200, ysize=200):
    with fits.open(fdata) as img:
        image = img[0].data.astype(np.float32)
        if fbias is not None:
            image = image - fits.getdata(fbias).astype(np.float32)
        if len(image.shape) == 3:
            image = image.mean(axis=0)  # merge color channels

        # cutout
        ny, nx = image.shape
        if xcent is None:
            xmin, xmax = 0, nx
        else:
            xmin, xmax = xcent - xsize // 2, xcent + xsize // 2
        if ycent is None:
            ymin, ymax = 0, ny
        else:
            ymin, ymax = ycent - ysize // 2, ycent + ysize // 2
        image = image[ymin:ymax, xmin:xmax]

        # detection source
        index = image > (np.max(image) * limit)
        index = binary_erosion(index)
        index = binary_dilation(index)
        #print(index)
        lab, n = label(index, return_num=True)
        print("spot number",n)
        if n == 0:
            return ['limit too high', ] * 9
        if n > 1:
            return ['limit too low', ] * 9

        clip = sigmaclip(image[~index], 3, 3)[0]
        bkg = clip.mean()
        rms = clip.std()
        #index = (image > bkg+rms*7)
        index = (image > (bkg + rms * 7)) * (image > (np.max(image) * limit))
        #index = (image>0)

        index = binary_erosion(index)
        index = binary_dilation(index)
        lab, n = label(index, return_num=True)
        if n == 0:
            return ['0', ] * 9
        if n > 1:
            return ['too many', ] * 9

        # position center
        ny, nx = image.shape
        yy, xx = np.mgrid[:ny, :nx]
        x0 = xx[index].mean()
        y0 = yy[index].mean()
        fmax = image[index].max()
        print("x0",x0,"y0",y0)
        # fit Gaussian
        index = binary_dilation(index, iterations=3)
        fitter = LevMarLSQFitter()
        model = Gaussian2D(x_mean=x0, y_mean=y0, amplitude=fmax)
        if sigma is not None:
            model.x_stddev = sigma
            model.y_stddev = sigma
        fit = fitter(model, xx[index], yy[index], image[index])
        im = np.zeros((20,20))
        for i in range(20):
            for j in range(20):
                im[i,j] = fit.evaluate(xx[int(y0)-10+i,int(x0)-10+j],yy[int(y0)-10+i,int(x0)-10+j],\
                    fit.amplitude, fit.x_mean, fit.y_mean, fit.x_stddev, fit.y_stddev, fit.theta)
        plt.subplot(131)
        plt.imshow(image[int(y0)-10:int(y0)+10,int(x0)-10:int(x0)+10])
        plt.subplot(132)
        plt.imshow(im)
        plt.subplot(133)
        plt.imshow(image[int(y0)-10:int(y0)+10,int(x0)-10:int(x0)+10]-im)
        plt.colorbar()
        plt.show()
        ellp = fit.y_stddev / fit.x_stddev
        if ellp > 1:
            ellp = 1 / ellp
        ellp = 1 - ellp
        print('x_size = ', fit.x_stddev.value * 2.355,'\ty_size = ', fit.y_stddev.value * 2.355,'\tellp = ', ellp )
        
        return ["{:.2f}".format(x0 + xmin), "{:.2f}".format(y0 + ymin),
                "{:.2f}".format(fit.x_mean.value + xmin), "{:.2f}".format(fit.y_mean.value + ymin),
                "{:.3f}".format(fit.x_stddev.value), "{:.3f}".format(fit.y_stddev.value),
                "{:.3f}".format(fit.x_stddev.value * 2.355), "{:.3f}".format(fit.y_stddev.value * 2.355),
                "{:.5f}".format(ellp)]


def _write_log(flog, fname, result):
    if os.path.exists(flog) is False:
        with open(flog, 'w') as f:
            f.write('# name, x, y, x_fit, y_fit, x_sigma, y_sigma, x_fwhm, y_fwhm, ellipticity\n')
    with open(flog, 'a') as f:
        f.write(fname + ', ' + ', '.join(result) + '\n')


def _run(flist, fbias, limit, sigma, xcent, ycent, xsize, ysize, flog):
    for f in flist:
        res = measure_beam(f, fbias=fbias, limit=limit, sigma=sigma,
                           xcent=xcent, ycent=ycent, xsize=xsize, ysize=ysize)
        if flog is not None:
            _write_log(flog, f, res)


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('flist', type=str, nargs='*')
    parser.add_argument('--bias', default=None, type=str,
                        help='bias file name')
    parser.add_argument('--limit', default=0.1, type=float,
                        help='the fraction of highest pixel value as lower limit for source detection')
    parser.add_argument('--sigma', default=None, type=float,
                        help='the initial sigma value for Gaussian fitting. If not set, ignore it.')
    parser.add_argument('--xcent', '-x', default=None, type=int,
                        help='x position of cutout center. If not set, use the full image')
    parser.add_argument('--ycent', '-y', default=None, type=int,
                        help='y position of cutout center. If not set, use the full image')
    parser.add_argument('--xsize', default=200, type=int,
                        help='x size of cutout, default is 200')
    parser.add_argument('--ysize', default=200, type=int,
                        help='y size of cutout, default is 200')
    parser.add_argument('--log', default=None, type=str,
                        help='output result table file')
    pars = parser.parse_args()
    _run(pars.flist, pars.bias, pars.limit, pars.sigma,
         pars.xcent, pars.ycent, pars.xsize, pars.ysize, pars.log)
